Skip to content

feat(auth): [aiohttp] Add mTLS reconfiguration logic when certificate mismatch for existing credentials & Agent Identity workloads - #18224

Open
agrawalradhika-cell wants to merge 79 commits into
mainfrom
cert-rotation-aiohttp
Open

feat(auth): [aiohttp] Add mTLS reconfiguration logic when certificate mismatch for existing credentials & Agent Identity workloads #18224
agrawalradhika-cell wants to merge 79 commits into
mainfrom
cert-rotation-aiohttp

Conversation

@agrawalradhika-cell

@agrawalradhika-cell agrawalradhika-cell commented Aug 26, 2026

Copy link
Copy Markdown
Contributor

feat: [aiohttp] Add mTLS reconfiguration logic when certificate mismatch for existing credentials & Agent Identity workloads

Changes included:

  • Adds mTLS reconfiguration logic when a certificate mismatch occurs.
  • Aligns async session behavior with the synchronous session by triggering a credential refresh and request retry on all 401 Unauthorized responses (not just mTLS).
  • Make sure to open an issue as a bug/issue before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea
  • Ensure the tests and linter pass
  • Code coverage does not decrease (if any source code was changed)
  • Appropriate docs were updated (if necessary)

Fixes #18227 #18227 🦕

@agrawalradhika-cell
agrawalradhika-cell requested review from a team as code owners August 26, 2026 04:00

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces client certificate rotation handling for asynchronous authorized sessions when encountering an unauthorized response under mTLS. The review feedback highlights a violation of the repository style guide regarding exception contract compliance, suggesting that the certificate parameter check should be wrapped in a try-except block to gracefully fall back to the original response rather than crashing. Additionally, the feedback recommends updating the corresponding unit tests to assert this resilient fallback behavior.

Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/tests/transport/aio/test_sessions_mtls.py Outdated
@agrawalradhika-cell agrawalradhika-cell changed the title feat: [aiohttp] Add reconfiguration logic when certificate mismatch for existing credentials & Agent Identity workloads feat: [aiohttp] Add mTLS reconfiguration logic when certificate mismatch for existing credentials & Agent Identity workloads Aug 26, 2026
agrawalradhika-cell and others added 2 commits August 26, 2026 10:34
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Handle exceptions during mTLS reconfiguration with warnings instead of errors.
…logs

Updated test logic to assert response instead of expecting an error.
…sync executor

Refactor unauthorized response handling to use async executor for MTLS parameter checks.
chore: Reset mTLS init task upon client certificate change
Signed-off-by: Radhika Agrawal <agrawalradhika@google.com>
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/tests/transport/aio/test_sessions_mtls.py
Comment thread packages/google-auth/tests/transport/aio/test_sessions_mtls.py Outdated
Comment thread packages/google-auth/tests/transport/aio/test_sessions_mtls.py
…eck after 401 check

chore: Refactor mTLS channel reconfiguration logic for adding mTLS check after 401 check
Implement mTLS rotation lock to prevent race conditions during certificate reconfiguration.
chore: Change warning to error log for mTLS channel reconfiguration failure.
chore: Refactor mTLS handling for unauthorized responses
Remove unnecessary continue statement after mTLS configuration.
Refactor tests for certificate rotation and error handling in AsyncAuthorizedSession. Update test names for clarity and ensure proper logging of errors.
Signed-off-by: Radhika Agrawal <agrawalradhika@google.com>
Signed-off-by: Radhika Agrawal <agrawalradhika@google.com>
Handle RefreshError during credential refresh to prevent unhandled exceptions.
Signed-off-by: Radhika Agrawal <agrawalradhika@google.com>
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/tests/transport/aio/test_sessions_mtls.py Outdated
Comment thread packages/google-auth/tests/transport/aio/test_sessions_mtls.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/tests/transport/aio/test_sessions_mtls.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Added refresh lock and counter to manage concurrent credential refreshes.
Limit the number of old authentication requests to 2 and ensure proper closure of the oldest requests.
Added tests for handling 401 responses with timeout and cancellation scenarios.
Signed-off-by: Radhika Agrawal <agrawalradhika@google.com>
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/mtls.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/mtls.py
Comment thread packages/google-auth/google/auth/aio/transport/mtls.py Outdated
Comment thread packages/google-auth/tests/transport/aio/test_sessions_mtls.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/sessions.py Outdated
Comment thread packages/google-auth/google/auth/aio/transport/mtls.py Outdated
Comment thread packages/google-auth/tests/transport/aio/test_mtls.py
Updated the 'check_parameters_for_unauthorized_response' function to include an optional client_cert_callback parameter and added detailed docstring for better understanding.
Store refresh counter at error for better tracking.
Updated mock patches to ensure correct assertions and added bound mocks for certificate parsing and caching.
Update mock authentication response handling in tests.
Signed-off-by: Radhika Agrawal <agrawalradhika@google.com>
else:
try:
await self._credentials.refresh(self._auth_request)
except NotImplementedError:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In _recover_auth_state(), catching NotImplementedError logs a debug message without returning response, allowing execution to fall through to return None (which signals successful recovery). This causes request() to dispatch two redundant retries with the unchanged credentials before finally returning the 401 response.

Consider returning response directly inside the except NotImplementedError: block (matching the RefreshError handling below):

                                except NotImplementedError:
                                    _LOGGER.debug(
                                        "Credentials do not implement refresh()."
                                    )
                                    return response

"""
call_cert_bytes, call_key_bytes = call_client_cert_callback()
if client_cert_callback:
call_cert_bytes, call_key_bytes = client_cert_callback()

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: AsyncAuthorizedSession now delegates mTLS parameter checking to google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response (which already handles missing certificates at lines 203–204). Because synchronous transports do not pass client_cert_callback, adding this parameter to _mtls_helper.py is unused. Consider reverting the changes to _mtls_helper.py to eliminate dead code and keep this PR focused on the aio package.

cached_cert
)
else:
cached_fingerprint = current_fingerprint

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When cached_cert is unset or falsy, assigning cached_fingerprint = current_fingerprint causes cached_fingerprint != current_cert_fingerprint in sessions.py to evaluate to False. If a session is initialized with a callback that initially yields (None, None), subsequent 401 recovery on an mTLS endpoint will detect no fingerprint change once certificates become available, skipping channel reconfiguration.

Consider assigning cached_fingerprint = None when cached_cert is falsy so the mismatch check triggers reconfiguration:

        if cached_cert:
            cached_fingerprint = _agent_identity_utils.get_cached_cert_fingerprint(
                cached_cert
            )
        else:
            cached_fingerprint = None
        return cached_fingerprint, current_cert_fingerprint

"Client certificate has changed, reconfiguring mTLS "
"channel."
)
if (

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: _recover_auth_state() only clears self._mtls_init_task when done() is true. If configure_mtls_channel() was called concurrently and is still pending, the call below returns the existing in-flight task without installing the new certificate callback, resulting in a redundant 401 retry. Consider awaiting self._mtls_init_task if it is still running before resetting it to None:

if self._mtls_init_task is not None:
    if not self._mtls_init_task.done():
        try:
            await self._mtls_init_task
        except Exception:
            pass
    self._mtls_init_task = None
await self.configure_mtls_channel(
    lambda: (call_cert_bytes, call_key_bytes)
)

"Failed to check client certificate parameters: %s. Proceeding with original response.",
e,
)
return response

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Returning response at line 414 inside the certificate inspection except block exits _recover_auth_state() before reaching await self._credentials.refresh(self._auth_request) at line 470.

When a 401 is caused by token expiration or revocation rather than cert rotation, skipping refresh breaks caller retries: because before_request checks only local timestamps against self.expiry, it keeps re-attaching the rejected token on subsequent calls. A transient inspection error (such as a disk I/O glitch or callback failure) therefore leaves the session stuck failing with 401.

Note that moving self._mtls_check_counter += 1 to a finally block would advance the counter on MutualTLSChannelError, causing concurrent tasks to swallow reconfiguration failures.

To fix this, consider removing return response so execution falls through to credential refresh, while keeping self._mtls_check_counter += 1 in the else block at line 458.

), mock.patch(
"aiohttp.ClientSession"
) as mock_session:
"""Tests that an exception in old_auth_request.close() does not abort configuration."""

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: In configure_mtls_channel(), old transports are buffered until len(self._old_auth_requests) >= 2 before closing. Because a single configuration call appends session._auth_request without calling close(), the eviction loop and exception suppression block in sessions.py remain unexercised here (the mock error is only caught at teardown in session.close()). Consider driving configure_mtls_channel() three times (clearing _mtls_init_task = None between calls) or pre-populating _old_auth_requests so the test exercises the eviction loop and verifies that close() exceptions do not abort configuration.

@@ -204,12 +224,18 @@ async def _do_configure():

old_auth_request = self._auth_request
self._auth_request = AiohttpRequest(session=new_session)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In configure_mtls_channel(), self._auth_request is reassigned to the new AiohttpRequest before the eviction loop, while old_auth_request is only appended to self._old_auth_requests after the loop completes. If the task is cancelled while awaiting oldest_auth_request.close(), old_auth_request is never tracked in self._old_auth_requests, causing session.close() to skip closing its underlying session and connection pool.

Popping oldest_auth_request before awaiting its close also risks leaving the in-eviction transport unclosed if cancellation occurs during the await.

Consider appending old_auth_request immediately, updating the eviction threshold to > 2, and popping oldest_auth_request only after its close completes:

                            old_auth_request = self._auth_request
                            self._auth_request = AiohttpRequest(session=new_session)
                            self._old_auth_requests.append(old_auth_request)
                            while len(self._old_auth_requests) > 2:
                                oldest_auth_request = self._old_auth_requests[0]
                                try:
                                    if hasattr(oldest_auth_request, "close"):
                                        res = oldest_auth_request.close()
                                        if inspect.isawaitable(res):
                                            await res
                                except Exception:
                                    pass
                                self._old_auth_requests.pop(0)

)

with pytest.raises(TimeoutError):
await session.request("GET", "https://example.com", max_allowed_time=0.01)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: max_allowed_time=0.01 (10ms) against real clock time risks flaking on loaded CI runners if event loop scheduling latency exceeds ~8.5ms during initial mock dispatch, which aborts before 401 recovery is entered and leaves mock_resp_401 unclosed. Consider setting max_allowed_time=0.1 (or 0.2) to provide ample headroom for mock dispatch while keeping unit test runtime brisk and reliably timing out inside slow_refresh (10s):

        with pytest.raises(TimeoutError):
            await session.request("GET", "https://example.com", max_allowed_time=0.1)


assert resp == mock_resp_200
mock_check.assert_called_once()
mock_conf.assert_called_once_with(mock.ANY)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: In test_cert_rotation_success_and_retry and test_psc_endpoint_triggers_cert_rotation, asserting mock_conf.assert_called_once_with(mock.ANY) accepts any argument and does not verify that the newly fetched certificates are forwarded into configure_mtls_channel(). Consider capturing the callback and asserting its return value to ensure rotated credentials are sent:

            mock_conf.assert_called_once()
            cb = (
                mock_conf.call_args.args[0]
                if mock_conf.call_args.args
                else mock_conf.call_args.kwargs["client_cert_callback"]
            )
            assert cb() == (new_cert, new_key)

optional callback which returns client certificate bytes and private
key bytes both in PEM format.

Returns:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: The docstring lists only non-None types for the 4-tuple return, but lines 203–204 return (None, None, None, None) when mTLS is inactive or certificates are absent. Consider documenting the sentinel return and using Optional types in the docstring to keep the return contract explicit:

    Returns:
        Tuple[Optional[bytes], Optional[bytes], Optional[str], Optional[str]]:
            is_mtls, current_cert_fingerprint, call_cert_bytes, call_key_bytes.
            Returns (None, None, None, None) if mTLS is disabled or no client certificate is present.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add cert rotation handling for aiohttp (Async HTTP)

6 participants